{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torch.utils.data as data\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.datasets import MNIST\n",
    "from datasets import MNIST_truncated\n",
    "from utils import AddGaussianNoise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "60000\n",
      "tensor([[  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
      "           0,  51, 159, 253, 159,  50,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
      "          48, 238, 252, 252, 252, 237,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,  54,\n",
      "         227, 253, 252, 239, 233, 252,  57,   6,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,  10,  60, 224,\n",
      "         252, 253, 252, 202,  84, 252, 253, 122,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0, 163, 252, 252,\n",
      "         252, 253, 252, 252,  96, 189, 253, 167,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,  51, 238, 253, 253,\n",
      "         190, 114, 253, 228,  47,  79, 255, 168,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,  48, 238, 252, 252, 179,\n",
      "          12,  75, 121,  21,   0,   0, 253, 243,  50,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,  38, 165, 253, 233, 208,  84,\n",
      "           0,   0,   0,   0,   0,   0, 253, 252, 165,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   7, 178, 252, 240,  71,  19,  28,\n",
      "           0,   0,   0,   0,   0,   0, 253, 252, 195,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,  57, 252, 252,  63,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0, 253, 252, 195,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0, 198, 253, 190,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0, 255, 253, 196,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,  76, 246, 252, 112,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0, 253, 252, 148,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,  85, 252, 230,  25,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   7, 135, 253, 186,  12,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,  85, 252, 223,   0,   0,   0,   0,   0,\n",
      "           0,   0,   0,   7, 131, 252, 225,  71,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,  85, 252, 145,   0,   0,   0,   0,   0,\n",
      "           0,   0,  48, 165, 252, 173,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,  86, 253, 225,   0,   0,   0,   0,   0,\n",
      "           0, 114, 238, 253, 162,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,  85, 252, 249, 146,  48,  29,  85, 178,\n",
      "         225, 253, 223, 167,  56,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,  85, 252, 252, 252, 229, 215, 252, 252,\n",
      "         252, 196, 130,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,  28, 199, 252, 252, 253, 252, 252, 233,\n",
      "         145,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,  25, 128, 252, 253, 252, 141,  37,\n",
      "           0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0],\n",
      "        [  0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,\n",
      "           0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0,   0]],\n",
      "       dtype=torch.uint8)\n",
      "torch.Size([60000, 28, 28])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(<torch.utils.data.dataloader.DataLoader at 0x2c211dd2630>,\n",
       " <torch.utils.data.dataloader.DataLoader at 0x2c23a768128>,\n",
       " <datasets.MNIST_truncated at 0x2c23a768da0>,\n",
       " <datasets.MNIST_truncated at 0x2c211dd2e48>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_dataloader(datadir, train_bs, test_bs, dataidxs=None, noise_level=0, net_id=None, total=0):\n",
    "\n",
    "    dl_obj = MNIST_truncated\n",
    "\n",
    "    transform_train = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        AddGaussianNoise(0., noise_level, net_id, total)])\n",
    "\n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        AddGaussianNoise(0., noise_level, net_id, total)])\n",
    "\n",
    "    train_ds = dl_obj(datadir, dataidxs=dataidxs, train=True, transform=transform_train, download=True)\n",
    "    test_ds = dl_obj(datadir, train=False, transform=transform_test, download=True)\n",
    "    \n",
    "    train_ds.data\n",
    "\n",
    "    # img, _ = train_ds.__getitem__(0)\n",
    "    # img.show()\n",
    "\n",
    "    train_dl = data.DataLoader(dataset=train_ds, batch_size=train_bs, shuffle=True, drop_last=False)\n",
    "    test_dl = data.DataLoader(dataset=test_ds, batch_size=test_bs, shuffle=False, drop_last=False)\n",
    "\n",
    "    print(len(train_ds))\n",
    "    print(train_ds.data.shape)\n",
    "    \n",
    "    return train_dl, test_dl, train_ds, test_ds\n",
    "\n",
    "# get_dataloader(\"./data/\", 32, 32, dataidxs=None, noise_level=0.01)\n",
    "get_dataloader(\"./data/\", 32, 32)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "FedAvg",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
